import re
import random
import json


def remove_comments(code: str) -> str:
    """
    Removes comments from Java, JavaScript, and Python code.
    
    Args:
        code (str): The input code string.
        
    Returns:
        str: The code string with comments removed.
    """
    # Regex for single-line comments and multi-line comments
    java_js_single_line = r'//.*?$'  # Single-line comments for Java/JavaScript
    python_single_line = r'#.*?$'   # Single-line comments for Python
    multi_line = r'/\*.*?\*/'       # Multi-line comments for Java/JavaScript
    
    # Combine regex patterns
    combined_pattern = f"({java_js_single_line}|{python_single_line}|{multi_line})"
    
    # Remove comments using regex with multiline and dotall flags
    no_comments = re.sub(combined_pattern, '', code, flags=re.MULTILINE | re.DOTALL)
    return no_comments



def split_code(code, min_mid_len=10, max_mid_len=200):
    code_len = len(code)
    
    # 生成随机的mid长度，保证在10到200之间
    mid_len = random.randint(min_mid_len, max_mid_len)
    start_idx = random.randint(0, code_len - mid_len)
    end_idx = start_idx + mid_len

    # 使用正则表达式来匹配变量名
    # 假设变量名符合Python命名规范，以字母、数字、下划线组成，并不以数字开头
    var_pattern = r'\b\w+\b'
    
    # 确保start_idx和end_idx不在变量名的中间
    def adjust_idx(idx, direction):
        while idx > 0 and idx < code_len:
            # 找到当前索引处的字符是否为变量名的一部分
            if re.match(var_pattern, code[idx:]):
                if direction == 'left':
                    idx -= 1
                elif direction == 'right':
                    idx += 1
            else:
                break
        return idx

    # 调整start_idx和end_idx，避免截断变量名
    start_idx = adjust_idx(start_idx, 'left')
    end_idx = adjust_idx(end_idx, 'right')

    # 提取 pre, mid 和 suf
    pre = code[:start_idx]
    mid = code[start_idx:end_idx]
    suf = code[end_idx:]

    pre.replace('<FILL_ME>', '')
    suf.replace('<FILL_ME>', '')

    return {
        "instruction": f'{pre}<FILL_ME>{suf}',
        "input": "",
        "output": mid,
        "system": "",
        "history": []
        }

def read_json_to_list(filename):
    with open(filename, 'r')as f:
        return json.load(f)
    
def write_json_to_file(filename, content):
    with open(filename, 'w')as f:
        return json.dump(content, f, ensure_ascii=False, indent=4)


if __name__ == '__main__':
    code_type = ['py', 'js', 'java']

    for type in code_type:
        print(len(read_json_to_list(f'/data/xingshiyou-slurm/codeLLM/data/original_code_data/output_by_keyword_{type}.json')))

        split_jsonl = [split_code(remove_comments(code_content), 25, 500) for code_content in read_json_to_list(f'/data/xingshiyou-slurm/codeLLM/data/original_code_data/output_by_keyword_{type}.json') if '<FILL_ME>' not in code_content and len(remove_comments(code_content)) < 9950 and len(remove_comments(code_content)) > 500]

        print('总长度：', len(split_jsonl))

        output_file = f'/data/xingshiyou-slurm/codeLLM/data/train_data/train_{type}_by_keyword_10000.json'

        write_json_to_file(output_file, split_jsonl)

    